import sys
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QTreeView, QVBoxLayout, QWidget, QLabel, QHBoxLayout
from PyQt5.QtGui import QStandardItemModel, QStandardItem

# 示例数据
RESULT = {
    "NODE_RESULT": {
        "U1": np.array([1, 4, 7]),
        "U2": np.array([2, 5, 8]),
        "U3": np.array([3, 6, 9]),
        "Umag": np.sqrt(np.array([1, 4, 7]) ** 2 + np.array([2, 5, 8]) ** 2 + np.array([3, 6, 9]) ** 2),
        "RF1": np.array([10, 13, 16]),
        "RF2": np.array([11, 14, 17]),
        "RF3": np.array([12, 15, 18]),
        "RFmag": np.sqrt(np.array([10, 13, 16]) ** 2 + np.array([11, 14, 17]) ** 2 + np.array([12, 15, 18]) ** 2),
    },
    "ELEMENT_RESULT": {},
    "INTEGRATION_POINT_RESULT": {}
}

class ExampleApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle("QTreeView 示例")
        self.setGeometry(100, 100, 600, 400)

        # 创建 QTreeView 控件
        self.treeView = QTreeView(self)
        self.treeView.setHeaderHidden(True)  # 隐藏表头

        # 创建模型
        self.model = QStandardItemModel()
        self.model.setHorizontalHeaderLabels(["名称"])

        # 构建树形结构
        self.buildTree(RESULT, self.model.invisibleRootItem())

        # 设置模型
        self.treeView.setModel(self.model)

        # 创建一个标签用于显示选中节点的数据
        self.dataLabel = QLabel("请选择一个节点以查看数据", self)
        self.dataLabel.setWordWrap(True)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.treeView)
        layout.addWidget(self.dataLabel)

        # 创建一个中心窗口
        centralWidget = QWidget(self)
        centralWidget.setLayout(layout)
        self.setCentralWidget(centralWidget)

        # 连接信号与槽
        self.treeView.selectionModel().selectionChanged.connect(self.onSelectionChanged)

    def buildTree(self, data, parentItem):
        for key, value in data.items():
            item = QStandardItem(key)
            item.setEditable(False)
            item.setData(value)  # 将数据存储在自定义属性中
            parentItem.appendRow(item)
            if isinstance(value, dict):
                self.buildTree(value, item)

    def onSelectionChanged(self, selected):
        # 获取选中的节点
        print(f"selected:{selected}")
        indexes = selected.indexes()
        if indexes:
            index = indexes[0]
            print(f"index:{index}")
            item = self.model.itemFromIndex(index)
            print(f"data{item}")
            data = item.data()  # 获取存储的数据
            if isinstance(data, np.ndarray):
                self.dataLabel.setText(f"数据:{data}")
            elif isinstance(data, dict):
                self.dataLabel.setText(f"数据:{data.keys()}")
            else:
                self.dataLabel.setText(f"数据:{data}")
        else:
            self.dataLabel.setText("请选择一个节点以查看数据")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    ex = ExampleApp()
    ex.show()
    sys.exit(app.exec_())